# coding: utf-8
# Definitions of Multiple Traveling Salesman Problem with Visiting Constraints
import os
from random import sample, randint
import shutil

import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)

ACM_Blue = '#0182ac'
ACM_Yellow = '#ffd600'
ACM_Orange = '#fc9200'
ACM_Red = '#e33532'
ACM_Light_Blue = '#90cef2'
ACM_Green = '#adbf00'
ACM_Purple = '#78006b'
ACM_Dark_Blue = '#184c8c'
city_colors = ['C0', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9', 'C1']
line_colors = [ACM_Blue, ACM_Red, ACM_Green, ACM_Orange, ACM_Purple, ACM_Yellow, ACM_Dark_Blue, 'k', 'C9', 'C1',
               ACM_Blue, ACM_Red, ACM_Green, ACM_Orange, ACM_Purple, ACM_Yellow, ACM_Dark_Blue, 'k', 'C9', 'C1']
gradient_red = ['#ff0000', '#ff3300', '#ff6600', '#ff9900', '#ffcc00', '#ffff00']
gradient_red = gradient_red[::-1]


def make_dir(path):
    path = path.strip()  # remove space
    path = path.rstrip("\\")  # remove \
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)
        return True
    else:
        # print(path + ' exists.')
        return False


class Generator:
    def __init__(self, m: int, path, save_dir, depot=None):
        self.m = m  # number of salesmen
        self.file_dir, self.name = os.path.split(path)  # instance name
        self.data = np.loadtxt(path, skiprows=6, encoding='utf-8')
        self.data = np.insert(self.data, obj=0, values=0, axis=0)  # raw tsp data including 0
        self.cities = self.data[1:, 0].astype(int)  # city indexes
        self.n = len(self.cities)  # number of cities
        self.x, self.y = self.data[:, 1].astype(float), self.data[:, 2].astype(float)  # coordinates
        self.depot = depot  # depot
        # set accessibility matrix
        self.A = np.zeros(shape=(self.n + 1, self.m), dtype=int)
        self.A[self.depot, :] = 1  # depot accessible to all salesmen
        cities = self.cities[:].tolist()
        if self.depot:
            cities.remove(self.depot)
        exclusive = sample(cities, k=m)  # exclusive cities
        for m, city in enumerate(exclusive):
            self.A[city, m] = 1
        for i in range(1, len(self.A)):
            if self.A[i].sum() == 0:  # random access
                for j in range(m + 1):
                    self.A[i][j] = randint(0, 1)
            if self.A[i].sum() == 0:  # keep no missing cities
                j = randint(0, m - 1)
                self.A[i][j] = 1
        # save A
        filename, extension = os.path.splitext(self.name)
        save_dir = save_dir + rf'/{filename}/'
        make_dir(save_dir)
        shutil.copy(path, save_dir)
        save_path = save_dir + rf'A{self.m}.txt'
        np.savetxt(save_path, self.A.reshape(self.n + 1, self.m), fmt='%s', encoding='utf-8')


class Instance:
    def __init__(self, file, m: int):
        self.m = m  # number of salesmen
        self.file_dir, self.name = os.path.split(file)  # instance name
        self.name = self.name.split('.')[0]
        self.name = f'{self.name}_m{self.m}'
        self.data = np.loadtxt(file, skiprows=6, encoding='utf-8')
        self.data = np.insert(self.data, obj=0, values=0, axis=0)  # raw tsp data
        self.cities = self.data[1:, 0].astype(int).tolist()  # city indexes
        self.n = len(self.cities)  # number of cities
        self.x, self.y = self.data[:, 1].astype(float), self.data[:, 2].astype(float)  # coordinates
        self.distance_matrix = np.sqrt(np.square(self.x - self.x[np.arange(self.n + 1).reshape(self.n + 1, 1)])
                                       + np.square(self.y - self.y[np.arange(self.n + 1).reshape(self.n + 1, 1)]))
        self.distance_matrix[0, :], self.distance_matrix[:, 0] = 0, 0  # distance matrix
        # load A
        self.A = np.loadtxt(self.file_dir + rf'/A{m}.txt', encoding='utf-8')
        self.depot = 1  # depot
        # fast compute accessibility
        self.city_access = [np.where(self.A[city, :] == 1)[0].tolist() for city in [0] + self.cities]
        self.sale_access = [np.where(self.A[:, sale] == 1)[0].tolist() for sale in range(self.m)]
    
    def __repr__(self):
        return f'{self.name}, n={self.n}, m={self.m}, depot={self.depot}'
    
    def cost(self, route: list) -> float:
        return sum([self.distance_matrix[route[i - 1]][route[i]] for i in range(1, len(route))])
    
    def access(self, city: int = None, salesman: int = None):
        """ There can be only one of the params (city or salesman) can be filled.
        @param city: query the access of the city
        @param salesman: query the access of the salesman
        @return: a list of access (cities or salesmen) for the query
        """
        assert city is None or salesman is None
        if city is not None:
            return self.city_access[city]
        elif salesman is not None:
            return self.sale_access[salesman]
        else:
            raise (ValueError(f'city={city}, salesman={salesman}'))
    
    def draw(self, plot, routes=None, bold=None, text=False, dot_size=None, lw=3, rotate=False, dots=None, label=False):
        dot_size = round(np.log2(self.n)) * 1 if dot_size is None else dot_size
        # dot_size = 0
        x, y = (self.y, self.x) if rotate else (self.x, self.y)
        plot.xticks([])
        plot.yticks([])
        if routes is not None:
            alpha = np.array([0.2 for _ in range(len(routes))])
            if bold is not None:
                alpha[bold] = 1
            else:
                alpha[:] = 1
            for i, route in enumerate(routes):
                plot.plot(x[route], y[route], c=line_colors[i], lw=lw, alpha=alpha[i], label=f'sm{i + 1}')
        if text:
            for data in self.data[1:]:
                plot.text(data[1], data[2], s=data[0].astype(int), size=12)
        if label:
            if self.m == 2 or self.m == 4:
                plot.legend(loc=9, bbox_to_anchor=(0.5, 1.1), borderaxespad=0, ncol=4)
            elif self.m == 6:
                plot.legend(loc=9, bbox_to_anchor=(0.5, 1.18), borderaxespad=0, ncol=3)
            else:
                plot.legend(loc=9, bbox_to_anchor=(0.5, 1.18), borderaxespad=0, ncol=4)
        
        colors = ['#3e4145' for _ in range(self.n + 1)]
        if dots is not None:
            for dot in dots:
                colors[dot] = 'green'
        total_access = self.A.sum(axis=1).astype(int)
        dot_size = total_access * dot_size  # dot size reflected by total access
        plot.scatter(x[1:], y[1:], s=dot_size[1:] * 2, c=colors[1:])
        plot.scatter(x[self.depot], y[self.depot], marker='s', s=dot_size[1] * 2, c='k')
        return plot
    
    def draw_colorful_cities_with_label(self, plot, routes=None, bold=None, text=False, dot_size=None, rotate=False,
                                        dots=None, label=False):
        dot_size = round(np.log2(self.n)) * 1 if dot_size is None else dot_size
        # dot_size = 0
        x, y = (self.y, self.x) if rotate else (self.x, self.y)
        plot.xticks([])
        plot.yticks([])
        if routes is not None:
            alpha = np.array([0.2 for _ in range(len(routes))])
            if bold is not None:
                alpha[bold] = 1
            else:
                alpha[:] = 1
            for i, route in enumerate(routes):
                plot.plot(x[route], y[route], c=line_colors[i], lw=3, alpha=alpha[i], label=f'sm{i + 1}')
        if text:
            for data in self.data[1:]:
                plot.text(data[1], data[2], s=data[0].astype(int), size=12)
        colors = ['#3e4145' for _ in range(self.n + 1)]
        if dots is not None:
            for dot in dots:
                colors[dot] = 'green'
        total_access = self.A.sum(axis=1).astype(int)
        clusters = [[] for _ in range(self.m + 1)]
        for m in range(self.m + 1):
            for i in range(self.n + 1):
                if total_access[i] == m:
                    clusters[m].append([x[i], y[i]])
        for m in range(1, self.m + 1):
            cluster = np.array(clusters[m])
            plot.scatter(cluster[:, 0], cluster[:, 1], s=dot_size * m * 2, c=gradient_red[m], label=m)
        plot.scatter(x[self.depot], y[self.depot], marker='s', s=dot_size * 9, c='k')
        if label:
            if self.m == 2 or self.m == 4:
                plot.legend(loc=9, bbox_to_anchor=(0.5, 1.1), borderaxespad=0, ncol=4)
            elif self.m == 6:
                plot.legend(loc=9, bbox_to_anchor=(0.5, 1.18), borderaxespad=0, ncol=3)
            else:
                plot.legend(loc=9, bbox_to_anchor=(0.5, 1.18), borderaxespad=0, ncol=4)
        return plot
    
    def draw_2(self, plot, routes=None, bold=None, text=False, dot_size=None, rotate=False, dots=None):
        dot_size = round(np.log2(self.n)) * 2 if dot_size is None else dot_size
        dot_size *= 0
        x, y = (self.y, self.x) if rotate else (self.x, self.y)
        plot.scatter(x[1:], y[1:], s=dot_size, c='grey')
        if dots is not None:
            plot.scatter(x[dots], y[dots], s=dot_size, c=ACM_Red, marker='*')
        plot.xticks([])
        plot.yticks([])
        if routes is not None:
            alpha = np.array([0.2 for _ in range(len(routes))])
            alpha[bold] = 1
            for i, route in enumerate(routes):
                # plot.plot(x[route], y[route], c=line_colors[i], alpha=alpha[i], lw=4)
                plot.plot(x[route], y[route], c=line_colors[i], alpha=alpha[i], label=f'sm{i + 1}', lw=3)
                # plot.plot(x[route[:-1]], y[route[:-1]], c=line_colors[i], alpha=alpha[i], label=f'salesman{i+1}', lw=4)
                # plot.plot(x[route[-2:]], y[route[-2:]], c=line_colors[i], alpha=alpha[i], ls=':', lw=4)
        if text:
            # for data in self.data[routes[bold]]:
            #     plot.text(data[1], data[2], s=data[0].astype(int), size=dot_size / 3)
            for data in self.data[1:]:
                plot.text(data[1], data[2], s=data[0].astype(int), size=dot_size / 3)
        plot.scatter(x[self.depot], y[self.depot], marker='s', s=dot_size, c='k')
        
        if len(routes) == 2 or len(routes) == 4:
            plt.legend(ncol=4, loc='center', bbox_to_anchor=(0.5, 1.05))
        elif len(routes) == 6:
            plt.legend(ncol=3, loc='center', bbox_to_anchor=(0.5, 1.08))
        elif len(routes) == 8:
            plt.legend(ncol=4, loc='center', bbox_to_anchor=(0.5, 1.08))
        
        return plot
    
    def draw_3(self, plot, routes=None, bold=None, text=False, dot_size=None, rotate=False, dots=None):
        dot_size = round(np.log2(self.n)) * 2 if dot_size is None else dot_size
        dot_size = 0
        x, y = (self.y, self.x) if rotate else (self.x, self.y)
        plot.scatter(x[1:], y[1:], s=dot_size, c='grey')
        if dots is not None:
            plot.scatter(x[dots], y[dots], s=dot_size, c=ACM_Red, marker='*')
        plot.scatter(x[4], y[4], s=100, c='k', marker='s')
        plot.scatter(x[33], y[33], s=200, c=ACM_Red, marker='*')
        plot.xticks([])
        plot.yticks([])
        if routes is not None:
            alpha = np.array([0.2 for _ in range(len(routes))])
            alpha[bold] = 1
            for i, route in enumerate(routes):
                # plot.plot(x[route], y[route], c=line_colors[i], alpha=alpha[i])
                plot.plot(x[route[:-1]], y[route[:-1]], c=line_colors[i], alpha=alpha[i], label=f'salesman{i + 1}',
                          lw=4)
                plot.plot(x[route[-2:]], y[route[-2:]], c=line_colors[i], alpha=alpha[i], ls=':', lw=4)
        if text:
            # for data in self.data[routes[bold]]:
            #     plot.text(data[1], data[2], s=data[0].astype(int), size=dot_size / 3)
            for data in self.data[1:]:
                plot.text(data[1], data[2], s=data[0].astype(int), size=dot_size)
        
        plot.scatter(x[self.depot], y[self.depot], marker='s', s=dot_size, c='k')
        return plot


# set_I = ['att48', 'eil51', 'berlin51', 'eil76', 'eil101'], [2, 4, 6, 8]
# set_II = ['kroA150', 'kroA200'], [10, 20, 30, 40]
# set_III = ['gr431', 'gr666'], [40, 60, 80, 100]
#
# SET = set_III
# for case in SET[0]:
#     root = '../../data/TSPLIB/'
#     for m in SET[1]:
#         print(f'{case} , {m} ...')
#         file = case + '.tsp.txt'
#         instance = Generator(m=m, path=root + file, save_dir='../../data/TSPLIB_VC/')

# for root, dirs, files in os.walk('../../data/TSPLIB/'):
#     for file in files:
#         if file == 'pla33810.tsp.txt':
#             continue
#         for m in range(2, 42, 2):
#             print(f'{file} , {m} ...')
#             instance = Generator(m=m, path=root + file, save_dir='../../data/TSPLIB_VC/')


# instance = Instance(file=f'../../data/TSPLIB_VC/burma14.tsp/burma14.tsp.txt', m=4)
# print(instance.cities, instance.depot)
# print(instance.A)
# print(instance.city_access)
# print(instance.sale_access)
